from typing import Optional
import pytest
import torch
import torch_npu
from unified_attn import unified_attention
import test_common
import random
import numpy as np

# NUM_HEADS = [(4, 4)]
# HEAD_SIZES = [128]
# BLOCK_SIZES = [16]
# DTYPES = [torch.float16, torch.bfloat16]
# NUM_BLOCKS = [32]

DTYPES = [torch.float16, torch.bfloat16]
SEQ_LENS = [(1, 512), (1, 1024), (4, 512), (4, 1024)]
NUM_HEADS = [(8, 8), (16, 8), (32, 8), (32, 4)]
HEAD_SIZES = [64, 128, 256]
BLOCK_SIZES = [16, 32, 64]
SLIDING_WINDOW = [None, 512, 1024]
NUM_BLOCKS = [32, 64, 128]
SOFT_CAP = [None, 10.0, 20.0]

# 10.20 版本之后，可以编译成功 silding_window存在问题

def ref_paged_attn(
    query: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    query_lens: list[int],
    kv_lens: list[int],
    block_tables: torch.Tensor,
    scale: float,
    sliding_window: Optional[int] = None,
    soft_cap: Optional[float] = None,
) -> torch.Tensor:
    num_seqs = len(query_lens)
    block_tables = block_tables.cpu().numpy()
    _, block_size, num_kv_heads, head_size = key_cache.shape

    outputs: list[torch.Tensor] = []
    start_idx = 0
    for i in range(num_seqs):
        query_len = query_lens[i]
        kv_len = kv_lens[i]
        q = query[start_idx:start_idx + query_len]
        q *= scale

        num_kv_blocks = (kv_len + block_size - 1) // block_size
        block_indices = block_tables[i, :num_kv_blocks]

        k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
        k = k[:kv_len]
        v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
        v = v[:kv_len]

        if q.shape[1] != k.shape[1]:
            k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
            v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
        attn = torch.einsum("qhd,khd->hqk", q, k).float()
        empty_mask = torch.ones(query_len, kv_len)
        mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
        if sliding_window is not None:
            sliding_window_mask = torch.triu(empty_mask,
                                             diagonal=kv_len -
                                             (query_len + sliding_window) +
                                             1).bool().logical_not()
            mask |= sliding_window_mask
        if soft_cap is not None and soft_cap > 0:
            attn = soft_cap * torch.tanh(attn / soft_cap)
        attn.masked_fill_(mask, float("-inf"))
        attn = torch.softmax(attn, dim=-1).to(v.dtype)
        out = torch.einsum("hqk,khd->qhd", attn, v)

        outputs.append(out)
        start_idx += query_len
    return torch.cat(outputs, dim=0)


# TEST_CASES = [
#     # 基础测试用例（不同序列长度组合）
#     [[(1, 512)], (8, 8), 64, None, torch.float16, 16, None, 32],
#     [[(1, 1024)], (8, 8), 64, None, torch.float16, 16, None, 32],
#     [[(4, 512)], (8, 8), 64, None, torch.float16, 16, None, 32],
#     [[(4, 1024)], (8, 8), 64, None, torch.float16, 16, None, 32],
#     [[(8, 2048)], (8, 8), 64, None, torch.float16, 16, None, 64],
#     [[(16, 1024)], (8, 8), 64, None, torch.float16, 16, None, 64],
    
#     # 不同头数配置
#     # ===============
#     # FAILED
#     # [[(1, 512)], (16, 8), 64, None, torch.float16, 16, None, 32],
#     # [[(1, 512)], (32, 8), 64, None, torch.float16, 16, None, 32],
#     # [[(1, 512)], (32, 4), 64, None, torch.float16, 16, None, 32],
#     # [[(1, 1024)], (16, 8), 64, None, torch.float16, 16, None, 32],
#     # [[(1, 1024)], (32, 8), 64, None, torch.float16, 16, None, 32],
#     # [[(4, 512)], (16, 8), 64, None, torch.float16, 16, None, 32],
#     # [[(4, 512)], (32, 4), 64, None, torch.float16, 16, None, 32],
#     # ==================
    
#     # 不同头维度
#     [[(1, 512)], (8, 8), 128, None, torch.float16, 16, None, 32],
#     [[(1, 512)], (8, 8), 256, None, torch.float16, 16, None, 32],
#     [[(1, 1024)], (8, 8), 128, None, torch.float16, 16, None, 32],
#     [[(1, 1024)], (8, 8), 256, None, torch.float16, 16, None, 64],
#     [[(4, 512)], (8, 8), 128, None, torch.float16, 16, None, 32],
#     [[(4, 512)], (8, 8), 256, None, torch.float16, 16, None, 64],
    
#     # # 不同块大小
#     [[(1, 512)], (8, 8), 64, None, torch.float16, 32, None, 32],
#     [[(1, 512)], (8, 8), 64, None, torch.float16, 64, None, 32],
#     [[(1, 1024)], (8, 8), 64, None, torch.float16, 32, None, 32],
#     [[(1, 1024)], (8, 8), 64, None, torch.float16, 64, None, 32],
#     [[(4, 512)], (8, 8), 64, None, torch.float16, 32, None, 32],
#     [[(4, 512)], (8, 8), 64, None, torch.float16, 64, None, 32],
    
#     # # 滑动窗口配置
#     [[(1, 512)], (8, 8), 64, 512, torch.float16, 16, None, 32],
#     [[(1, 512)], (8, 8), 64, 1024, torch.float16, 16, None, 32],
#     [[(1, 1024)], (8, 8), 64, 512, torch.float16, 16, None, 32],
#     [[(1, 1024)], (8, 8), 64, 1024, torch.float16, 16, None, 32],
#     [[(4, 512)], (8, 8), 64, 512, torch.float16, 16, None, 32],
#     [[(4, 512)], (8, 8), 64, 1024, torch.float16, 16, None, 32],
    
#     # # 不同数据类型
#     [[(1, 512)], (8, 8), 64, None, torch.bfloat16, 16, None, 32],
#     [[(1, 1024)], (8, 8), 64, None, torch.bfloat16, 16, None, 32],
#     [[(4, 512)], (8, 8), 64, None, torch.bfloat16, 16, None, 32],
#     # =================
#     # [[(1, 512)], (16, 8), 64, None, torch.bfloat16, 16, None, 32],
#     # =================
#     [[(1, 512)], (8, 8), 128, None, torch.bfloat16, 16, None, 32],
#     [[(1, 512)], (8, 8), 64, 512, torch.bfloat16, 16, None, 32],
    
#     # # 软上限配置
#     [[(1, 512)], (8, 8), 64, None, torch.float16, 16, 10.0, 32],
#     [[(1, 512)], (8, 8), 64, None, torch.float16, 16, 20.0, 32],
#     [[(1, 1024)], (8, 8), 64, None, torch.float16, 16, 10.0, 32],
#     [[(1, 1024)], (8, 8), 64, None, torch.float16, 16, 20.0, 32],
#     [[(4, 512)], (8, 8), 64, None, torch.float16, 16, 10.0, 32],
#     [[(4, 512)], (8, 8), 64, None, torch.float16, 16, 20.0, 32],
    
#     # # 不同块数量
#     [[(1, 512)], (8, 8), 64, None, torch.float16, 16, None, 64],
#     [[(1, 512)], (8, 8), 64, None, torch.float16, 16, None, 128],
#     [[(1, 1024)], (8, 8), 64, None, torch.float16, 16, None, 64],
#     [[(1, 1024)], (8, 8), 64, None, torch.float16, 16, None, 128],
#     [[(4, 512)], (8, 8), 64, None, torch.float16, 16, None, 64],
#     [[(4, 512)], (8, 8), 64, None, torch.float16, 16, None, 128],
    
#     # 组合配置1：多头+大序列
#     # ================
#     # [[(1, 2048)], (32, 8), 128, 1024, torch.float16, 32, 10.0, 64],
#     # [[(1, 2048)], (32, 4), 128, 1024, torch.float16, 32, 10.0, 64],
#     # [[(2, 2048)], (32, 8), 128, 1024, torch.float16, 32, 10.0, 64],
#     # [[(1, 4096)], (32, 8), 128, 2048, torch.float16, 64, 20.0, 128],
#     # [[(1, 4096)], (16, 8), 128, 2048, torch.float16, 64, 20.0, 128],
#     # [[(2, 2048)], (16, 8), 128, 1024, torch.float16, 32, 10.0, 64],
#     # =================
    
#     # 组合配置2：滑动窗口+软上限
#     [[(1, 512)], (8, 8), 64, 512, torch.float16, 16, 10.0, 32],
#     [[(1, 512)], (8, 8), 64, 512, torch.float16, 16, 20.0, 32],
#     [[(1, 1024)], (8, 8), 64, 512, torch.float16, 16, 10.0, 32],
#     [[(1, 1024)], (8, 8), 64, 1024, torch.float16, 16, 20.0, 32],
#     # =================
#     # [[(4, 512)], (16, 8), 64, 512, torch.float16, 16, 10.0, 32],
#     # [[(4, 512)], (16, 8), 64, 1024, torch.float16, 16, 20.0, 32],
#     # =================
    
#     # 组合配置3：大头部尺寸+大区块
#     [[(1, 512)], (8, 8), 256, None, torch.float16, 64, None, 32],
#     [[(1, 1024)], (8, 8), 256, None, torch.float16, 64, None, 64],
#     # ================
#     # [[(1, 1024)], (16, 8), 256, None, torch.float16, 64, None, 64],
#     # ================

#     [[(4, 512)], (8, 8), 256, None, torch.float16, 64, None, 32],

#     # ================
#     # [[(4, 512)], (16, 8), 256, None, torch.float16, 64, None, 64],
#     # ================

#     [[(1, 2048)], (8, 8), 256, 1024, torch.float16, 64, 20.0, 128],
    
#     # # 组合配置4：float32+特殊参数
#     [[(1, 512)], (8, 8), 64, 512, torch.bfloat16, 16, 10.0, 32],
#     # =================
#     # [[(1, 512)], (16, 8), 64, 512, torch.bfloat16, 16, 10.0, 32],
#     # =================
#     [[(1, 1024)], (8, 8), 128, 1024, torch.bfloat16, 32, 20.0, 64],
#     [[(4, 512)], (8, 8), 128, 512, torch.bfloat16, 32, 10.0, 32],
#     # =================
#     [[(4, 512)], (32, 4), 128, 512, torch.bfloat16, 32, 10.0, 64],
#     [[(1, 2048)], (16, 8), 128, 1024, torch.bfloat16, 64, 20.0, 128],
    
#     # # 长序列特殊测试
#     [[(1, 8192)], (32, 8), 128, 2048, torch.float16, 64, 20.0, 128],
#     [[(1, 8192)], (32, 4), 128, 2048, torch.float16, 64, 20.0, 128],
#     [[(2, 4096)], (32, 8), 128, 2048, torch.float16, 64, 20.0, 128],
#     [[(1, 16384)], (32, 8), 128, 4096, torch.float16, 64, 20.0, 256],
#     [[(1, 16384)], (64, 8), 128, 4096, torch.float16, 64, 20.0, 256],
#     [[(2, 8192)], (32, 8), 128, 2048, torch.float16, 64, 20.0, 128],
#     # ================
    
#     # 小批量多序列
#     [[(16, 256)], (8, 8), 64, 256, torch.float16, 16, 10.0, 32],
#     [[(16, 512)], (8, 8), 64, 512, torch.float16, 32, 10.0, 64],
#     # =============
#     [[(32, 256)], (16, 8), 64, 256, torch.float16, 16, 10.0, 32],
#     [[(32, 512)], (16, 8), 64, 512, torch.float16, 32, 10.0, 64],
#     [[(64, 256)], (16, 8), 64, 256, torch.float16, 16, 10.0, 64],
#     [[(64, 128)], (32, 4), 64, 128, torch.float16, 16, 10.0, 32],
    
#     # # 极端参数组合
#     [[(1, 128)], (64, 8), 64, 128, torch.float16, 16, 10.0, 32],
#     [[(1, 128)], (64, 4), 64, 128, torch.float16, 16, 10.0, 32],
#     [[(1, 256)], (64, 8), 128, 256, torch.float16, 32, 20.0, 64],
#     [[(1, 256)], (128, 16), 64, 256, torch.float16, 32, 20.0, 64],
#     [[(4, 128)], (64, 8), 64, 128, torch.float16, 16, 10.0, 32],
#     [[(4, 256)], (64, 8), 128, 256, torch.float16, 32, 20.0, 64],
    
#     # 混合参数集1
#     [[(2, 1024)], (16, 8), 128, 512, torch.float16, 32, 10.0, 64],
#     [[(2, 1024)], (32, 8), 128, 512, torch.float16, 32, 10.0, 64],
#     ===============
#     [[(2, 2048)], (16, 8), 128, 1024, torch.float16, 64, 20.0, 128],
#     [[(4, 1024)], (16, 8), 128, 512, torch.float16, 32, 10.0, 64],
#     [[(4, 1024)], (32, 4), 128, 512, torch.float16, 32, 10.0, 64],
#     [[(8, 512)], (16, 8), 128, 512, torch.float16, 32, 10.0, 64],
    
#     # # 混合参数集2
#     [[(1, 512)], (8, 8), 64, 512, torch.bfloat16, 32, 10.0, 32],
#     [[(1, 1024)], (16, 8), 64, 1024, torch.bfloat16, 32, 20.0, 64],
#     [[(4, 512)], (16, 8), 128, 512, torch.bfloat16, 32, 10.0, 64],
#     [[(4, 1024)], (32, 8), 128, 1024, torch.bfloat16, 64, 20.0, 128],
#     [[(8, 512)], (32, 8), 128, 512, torch.bfloat16, 32, 10.0, 64],
#     [[(8, 1024)], (32, 4), 128, 1024, torch.bfloat16, 64, 20.0, 128],
    
#     # 大模型典型配置
#     # [[(1, 4096)], (32, 8), 128, 2048, torch.float16, 64, 20.0, 128],
#     # [[(1, 8192)], (32, 8), 128, 4096, torch.float16, 64, 20.0, 256],
#     # [[(1, 8192)], (64, 8), 128, 4096, torch.float16, 64, 20.0, 256],
#     # [[(1, 16384)], (64, 16), 128, 4096, torch.float16, 64, 20.0, 512],
#     # [[(2, 4096)], (32, 8), 128, 2048, torch.float16, 64, 20.0, 128],
#     # [[(2, 8192)], (64, 8), 128, 4096, torch.float16, 64, 20.0, 256],
# ]

TEST_CASES = [
    [[(1, 512), (1, 1024), (4, 512), (4, 1024)], [8, 2], 128, None, torch.float16, 16, 20.0, 32],
]

# @pytest.mark.parametrize("seq_lens", SEQ_LENS)
# @pytest.mark.parametrize("num_heads", NUM_HEADS)
# @pytest.mark.parametrize("head_size", HEAD_SIZES)
# @pytest.mark.parametrize("block_size", BLOCK_SIZES)
# @pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
# @pytest.mark.parametrize("dtype", DTYPES)
# @pytest.mark.parametrize("soft_cap", SOFT_CAP)
# @pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
@pytest.mark.parametrize("test_case", TEST_CASES)
@torch.inference_mode()
def test_triton_unified_attn(
    # seq_lens: list[tuple[int, int]],
    # num_heads: tuple[int, int],
    # head_size: int,
    # sliding_window: Optional[int],
    # dtype: torch.dtype,
    # block_size: int,
    # soft_cap: Optional[float],
    # num_blocks: int,
    test_case
) -> None:
    (seq_lens, num_heads, head_size, sliding_window,
     dtype, block_size, soft_cap, num_blocks) = test_case
    torch.set_default_device("npu")
    torch.npu.set_device(2)
    torch.manual_seed(42)
    np.random.seed(42) 
    random.seed(42)
    num_seqs = len(seq_lens)
    query_lens = [x[0] for x in seq_lens]
    kv_lens = [x[1] for x in seq_lens]
    num_query_heads = num_heads[0]
    num_kv_heads = num_heads[1]
    assert num_query_heads % num_kv_heads == 0
    max_query_len = max(query_lens)
    max_kv_len = max(kv_lens)
    window_size = ((sliding_window - 1, 0) if sliding_window is not None else
                   (-1, -1))
    scale = head_size**-0.5

    query = torch.randn(sum(query_lens),
                        num_query_heads,
                        head_size,
                        dtype=dtype)
    key_cache = torch.randn(num_blocks,
                            block_size,
                            num_kv_heads,
                            head_size,
                            dtype=dtype)
    value_cache = torch.randn_like(key_cache)
    cu_query_lens = torch.tensor([0] + query_lens,
                                 dtype=torch.int32).cumsum(dim=0,
                                                           dtype=torch.int32)
    kv_lens = torch.tensor(kv_lens, dtype=torch.int32)

    max_num_blocks_per_seq = (max_kv_len + block_size - 1) // block_size
    block_tables = torch.randint(0,
                                 num_blocks,
                                 (num_seqs, max_num_blocks_per_seq),
                                 dtype=torch.int32)

    output = torch.empty_like(query)

    unified_attention(
        q=query,
        k=key_cache,
        v=value_cache,
        out=output,
        cu_seqlens_q=cu_query_lens,
        seqused_k=kv_lens,
        max_seqlen_q=max_query_len,
        max_seqlen_k=max_kv_len,
        softmax_scale=scale,
        causal=True,
        window_size=window_size,
        block_table=block_tables,
        softcap=soft_cap if soft_cap is not None else 0,
    )

    ref_output = ref_paged_attn(
        query=query,
        key_cache=key_cache,
        value_cache=value_cache,
        query_lens=query_lens,
        kv_lens=kv_lens,
        block_tables=block_tables,
        scale=scale,
        sliding_window=sliding_window,
        soft_cap=soft_cap,
    )
    # atol, rtol = 1.5e-1, 1e-1
    # torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol), f"{torch.max(torch.abs(output - ref_output))}"
    sigtype_dict = {
        torch.float16: "float16",
        torch.bfloat16: "bfloat16",
    }
    sigtype = sigtype_dict[dtype]
    test_common.validate_cmp(sigtype, output, ref_output)